package crawler

import (
	"errors"
	"log"

	"github.com/boltdb/bolt"
)

type (
	//KVStore KV存储
	KVStore struct {
		DB *bolt.DB
	}
)

//OpenDb 打开数据库
func NewKVStore(dbfile string) *KVStore {
	db, err := bolt.Open("cache", 0600, nil)
	if err != nil {
		log.Fatal(err)
	}
	return &KVStore{DB: db}
}

//DestoryCache 关闭数据库
func (c *KVStore) CloseStore() {
	c.DB.Close()
}

//Put
func (c *KVStore) Put(db, key, value ByteContent) {
	var err error
	err = c.DB.Update(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(db))
		if b == nil {
			b, err = tx.CreateBucket([]byte(db))
			if err != nil {
				log.Println(err.Error())
			}
		}
		b.Put(key, value)
		return nil
	})
	if err != nil {
		log.Println(err.Error())
	}
}

//Get
func (c *KVStore) Get(db, key ByteContent) (ByteContent, error) {
	var err error
	var ret ByteContent
	c.DB.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(db))
		if b == nil {
			err = errors.New("数据表不存在")
		} else {
			ret = b.Get(key)
		}
		return nil
	})
	return ret, err
}
